{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#以下代码修改自naturomics的GitHub实现，包含三层CapsNet和后面的重构网络\n",
    "#改网络参数比较多，我们后面会只训练测试三层CapsNet。\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "epsilon = 1e-9\n",
    "batch_size = 8\n",
    "epoch = 1\n",
    "\n",
    "#margin loss 中调节上margin和下margind的权重\n",
    "lambda_val = 0.5\n",
    "#上margin与下margin的参数值\n",
    "m_plus = 0.9\n",
    "m_minus = 0.1\n",
    "\n",
    "# 路由更新c_ij所经过的迭代次数\n",
    "iter_routing = 3\n",
    "\n",
    "# Tensorboard 保存位置\n",
    "logdir ='logdir'\n",
    "# 数据集路径\n",
    "dataset_path = 'data/MNIST'\n",
    "is_training= True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 定义加载mnist的函数\n",
    "def load_mnist(path, is_training):\n",
    "\n",
    "    #trX将加载储存所有60000张灰度图\n",
    "    fd = open(os.path.join(path, 'train-images.idx3-ubyte'))\n",
    "    loaded = np.fromfile(file=fd, dtype=np.uint8)\n",
    "    trX = loaded[16:].reshape((60000, 28, 28, 1)).astype(np.float)\n",
    "\n",
    "    fd = open(os.path.join(path, 'train-labels.idx1-ubyte'))\n",
    "    loaded = np.fromfile(file=fd, dtype=np.uint8)\n",
    "    trY = loaded[8:].reshape((60000)).astype(np.float)\n",
    "\n",
    "    #teX将储存所有一万张测试用的图片\n",
    "    fd = open(os.path.join(path, 't10k-images.idx3-ubyte'))\n",
    "    loaded = np.fromfile(file=fd, dtype=np.uint8)\n",
    "    teX = loaded[16:].reshape((10000, 28, 28, 1)).astype(np.float)\n",
    "\n",
    "    fd = open(os.path.join(path, 't10k-labels.idx1-ubyte'))\n",
    "    loaded = np.fromfile(file=fd, dtype=np.uint8)\n",
    "    teY = loaded[8:].reshape((10000)).astype(np.float)\n",
    "\n",
    "    # 将所有训练图片表示为一个4维张量 [60000, 28, 28, 1]，其中每个像素值缩放到0和1之间\n",
    "    trX = tf.convert_to_tensor(trX / 255., tf.float32)\n",
    "\n",
    "    # one hot编码为 [num_samples, 10]\n",
    "    trY = tf.one_hot(trY, depth=10, axis=1, dtype=tf.float32)\n",
    "    teY = tf.one_hot(teY, depth=10, axis=1, dtype=tf.float32)\n",
    "\n",
    "    # 训练和测试时返回不同的数据\n",
    "    if is_training:\n",
    "        return trX, trY\n",
    "    else:\n",
    "        return teX / 255., teY\n",
    "\n",
    "def get_batch_data():\n",
    "    trX, trY = load_mnist(dataset_path, True)\n",
    "\n",
    "    # 每次产生一个切片\n",
    "    data_queues = tf.train.slice_input_producer([trX, trY])\n",
    "\n",
    "    # 对队列中的样本进行乱序处理\n",
    "    X, Y = tf.train.shuffle_batch(data_queues,\n",
    "                                  batch_size=batch_size,\n",
    "                                  capacity=batch_size * 64,\n",
    "                                  min_after_dequeue=batch_size * 32,\n",
    "                                  allow_smaller_final_batch=False)\n",
    "    return (X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 通过定义类和对象的方式定义Capssule层级\n",
    "class CapsLayer(object):\n",
    "    ''' Capsule layer 类别参数有：\n",
    "    Args:\n",
    "        input: 一个4维张量\n",
    "        num_outputs: 当前层的Capsule单元数量\n",
    "        vec_len: 一个Capsule输出向量的长度\n",
    "        layer_type: 选择'FC' 或 \"CONV\", 以确定是用全连接层还是卷积层\n",
    "        with_routing: 当前Capsule是否从较低层级中Routing而得出输出向量\n",
    "\n",
    "    Returns:\n",
    "        一个四维张量\n",
    "    '''\n",
    "    def __init__(self, num_outputs, vec_len, with_routing=True, layer_type='FC'):\n",
    "        self.num_outputs = num_outputs\n",
    "        self.vec_len = vec_len\n",
    "        self.with_routing = with_routing\n",
    "        self.layer_type = layer_type\n",
    "\n",
    "    def __call__(self, input, kernel_size=None, stride=None):\n",
    "        '''\n",
    "        当“Layer_type”选择的是“CONV”，我们将使用 'kernel_size' 和 'stride'\n",
    "        '''\n",
    "\n",
    "        # 开始构建卷积层\n",
    "        if self.layer_type == 'CONV':\n",
    "            self.kernel_size = kernel_size\n",
    "            self.stride = stride\n",
    "\n",
    "            # PrimaryCaps层没有Routing过程\n",
    "            if not self.with_routing:\n",
    "                # 卷积层为 PrimaryCaps 层（CapsNet第二层）, 并将第一层卷积的输出张量作为输入。\n",
    "                # 输入张量的维度为： [batch_size, 20, 20, 256]\n",
    "                assert input.get_shape() == [batch_size, 20, 20, 256]\n",
    "\n",
    "                # # 从CapsNet输出向量的每一个分量开始执行卷积，每个分量上执行带32个卷积核的9×9标准卷积\n",
    "                # capsules = []\n",
    "                # for i in range(self.vec_len):\n",
    "                #     # 所有Capsule的一个分量，其维度为: [batch_size, 6, 6, 32]，即6×6×1×32\n",
    "                #     with tf.variable_scope('ConvUnit_' + str(i)):\n",
    "                #         caps_i = tf.contrib.layers.conv2d(input, self.num_outputs,\n",
    "                #                                           self.kernel_size, self.stride,\n",
    "                #                                           padding=\"VALID\")\n",
    "                #\n",
    "                #         # 将一般卷积的结果张量拉平，并为添加到列表中\n",
    "                #         caps_i = tf.reshape(caps_i, shape=(batch_size, -1, 1, 1))\n",
    "                #         capsules.append(caps_i)\n",
    "                #\n",
    "                # # 为将卷积后张量各个分量合并为向量做准备\n",
    "                # assert capsules[0].get_shape() == [batch_size, 1152, 1, 1]\n",
    "                #\n",
    "                # # 合并为PrimaryCaps的输出张量，即6×6×32个长度为8的向量，合并后的维度为 [batch_size, 1152, 8, 1]\n",
    "                # capsules = tf.concat(capsules, axis=2)\n",
    "                # # 将每个Capsule 向量投入非线性函数squash进行缩放与激活,第二层输出的向量要经过缩放\n",
    "                # capsules = squash(capsules)\n",
    "                # assert capsules.get_shape() == [batch_size, 1152, 8, 1]\n",
    "                # return(capsules)\n",
    "\n",
    "                # 以下更新后的计算方法\n",
    "                capsules = tf.contrib.layers.conv2d(input, self.num_outputs * self.vec_len,\n",
    "                                                    self.kernel_size, self.stride, padding=\"VALID\")\n",
    "                capsules = tf.reshape(capsules, (batch_size, -1, self.vec_len, 1))\n",
    "\n",
    "                # [batch_size, 1152, 8, 1]\n",
    "                capsules = squash(capsules)\n",
    "                assert capsules.get_shape() == [batch_size, 1152, 8, 1]\n",
    "                return (capsules)\n",
    "\n",
    "        if self.layer_type == 'FC':\n",
    "\n",
    "            # DigitCaps 带有Routing过程\n",
    "            if self.with_routing:\n",
    "                # CapsNet 的第三层 DigitCaps 层是一个全连接网络\n",
    "                # 将输入张量重建为 [batch_size, 1152, 1, 8, 1]\n",
    "                self.input = tf.reshape(input, shape=(batch_size, -1, 1, input.shape[-2].value, 1))\n",
    "\n",
    "                with tf.variable_scope('routing'):\n",
    "                    # 初始化b_IJ的值为零，且维度满足: [1, 1, num_caps_l, num_caps_l_plus_1, 1]\n",
    "                    b_IJ = tf.constant(np.zeros([1, input.shape[1].value, self.num_outputs, 1, 1], dtype=np.float32))\n",
    "                    # 使用定义的Routing过程计算权值更新与s_j\n",
    "                    capsules = routing(self.input, b_IJ)\n",
    "                    # 将s_j投入 squeeze 函数以得出 DigitCaps 层的输出向量\n",
    "                    capsules = tf.squeeze(capsules, axis=1)\n",
    "\n",
    "            return(capsules)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 定义路由算法的过程\n",
    "def routing(input, b_IJ):\n",
    "    ''' 路由算法\n",
    "\n",
    "    Args:\n",
    "        input: 输入张量的维度为 [batch_size, num_caps_l=1152, 1, length(u_i)=8, 1]\n",
    "               其中num_caps_l为上一层（PrimaryCaps）的Capsule单元数量\n",
    "    Returns:\n",
    "        返回的张量维度为 [batch_size, num_caps_l_plus_1, length(v_j)=16, 1]\n",
    "        表征了i+1层的输出向量 `v_j`，num_caps_l_plus_1 为DigitCaps层的输出数\n",
    "    Notes:\n",
    "        u_i 表示l层中 capsule i 的输出向量\n",
    "        v_j 表示l+1层中 capsule j 的输出向量\n",
    "     '''\n",
    "\n",
    "    # 定义W的张量维度为 [num_caps_j, num_caps_i, len_u_i, len_v_j]\n",
    "    # W_ij共有1152×10个，每一个的维度为8×16\n",
    "    W = tf.get_variable('Weight', shape=(1, 1152, 10, 8, 16), dtype=tf.float32,\n",
    "                        initializer=tf.random_normal_initializer(stddev=0.01))\n",
    "\n",
    "    # 论文中的 Eq.2, 计算 u_hat\n",
    "    # 在使用 W 和u_i计算u_hat前，先调整张量维度\n",
    "    # input => [batch_size, 1152, 10, 8, 1]\n",
    "    # W => [batch_size, 1152, 10, 8, 16]\n",
    "    input = tf.tile(input, [1, 1, 10, 1, 1])\n",
    "    W = tf.tile(W, [batch_size, 1, 1, 1, 1])\n",
    "    assert input.get_shape() == [batch_size, 1152, 10, 8, 1]\n",
    "\n",
    "    # 因为[8, 16].T x [8, 1] => [16, 1]，所以矩阵乘法在最后得出的维度为 [batch_size, 1152, 10, 16, 1]\n",
    "    u_hat = tf.matmul(W, input, transpose_a=True)\n",
    "    assert u_hat.get_shape() == [batch_size, 1152, 10, 16, 1]\n",
    "\n",
    "    # 前面是扩展的线性组合，后面是路由的部分，以下开始迭代路由过程更新耦合系数\n",
    "    # 对应论文中伪代码的第三行\n",
    "    for r_iter in range(iter_routing):\n",
    "        with tf.variable_scope('iter_' + str(r_iter)):\n",
    "            # 原论文伪代码第四行，计算softmax(b_ij)\n",
    "            # => [1, 1152, 10, 1,1]\n",
    "            c_IJ = tf.nn.softmax(b_IJ, dim=3)\n",
    "            c_IJ = tf.tile(c_IJ, [batch_size, 1, 1, 1, 1])\n",
    "            assert c_IJ.get_shape() == [batch_size, 1152, 10, 1, 1]\n",
    "\n",
    "            # 原论文伪代码第五行，根据更新的c_ij计算s_j\n",
    "            # 先利用 c_IJ 给 u_hat 加权，即在后两个维度采用对应元素的乘积\n",
    "            # => [batch_size, 1152, 10, 16, 1]\n",
    "            s_J = tf.multiply(c_IJ, u_hat)\n",
    "            # 在第二个维度上求和, 产生的张量维度为 [batch_size, 1, 10, 16, 1]\n",
    "            s_J = tf.reduce_sum(s_J, axis=1, keep_dims=True)\n",
    "            assert s_J.get_shape() == [batch_size, 1, 10, 16, 1]\n",
    "\n",
    "            # 原论文伪代码的第六行\n",
    "            # 使用 Eq.1 计算squashing非线性函数\n",
    "            v_J = squash(s_J)\n",
    "            assert v_J.get_shape() == [batch_size, 1, 10, 16, 1]\n",
    "\n",
    "            # 原论文伪代码的第七行\n",
    "            # reshape & tile v_j from [batch_size ,1, 10, 16, 1] to [batch_size, 10, 1152, 16, 1]\n",
    "            # then matmul in the last tow dim: [16, 1].T x [16, 1] => [1, 1], reduce mean in the\n",
    "            # batch_size dim, resulting in [1, 1152, 10, 1, 1]\n",
    "            v_J_tiled = tf.tile(v_J, [1, 1152, 1, 1, 1])\n",
    "            u_produce_v = tf.matmul(u_hat, v_J_tiled, transpose_a=True)\n",
    "            assert u_produce_v.get_shape() == [batch_size, 1152, 10, 1, 1]\n",
    "            b_IJ += tf.reduce_sum(u_produce_v, axis=0, keep_dims=True)\n",
    "\n",
    "    return(v_J)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def squash(vector):\n",
    "    ''' 根据原论文中 Eq. 1 定义squashing函数\n",
    "    Args:\n",
    "        vector: 一个 5-D 张量，其维度是 [batch_size, 1, num_caps, vec_len, 1],\n",
    "    Returns:\n",
    "        返回一个 5-D 张量，其第四和第五个维度经过了该非线性函数据算\n",
    "    '''\n",
    "    vec_squared_norm = tf.reduce_sum(tf.square(vector), -2, keep_dims=True)\n",
    "    scalar_factor = vec_squared_norm / (1 + vec_squared_norm) / tf.sqrt(vec_squared_norm + epsilon)\n",
    "    vec_squashed = scalar_factor * vector  # element-wise\n",
    "    return(vec_squashed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 以下定义整个 CapsNet 的架构与正向传播过程\n",
    "class CapsNet():\n",
    "    def __init__(self, is_training=True):\n",
    "        self.graph = tf.Graph()\n",
    "        with self.graph.as_default():\n",
    "            if is_training:\n",
    "                # 获取一个批量的训练数据\n",
    "                self.X, self.Y = get_batch_data()\n",
    "\n",
    "                self.build_arch()\n",
    "                self.loss()\n",
    "\n",
    "                # t_vars = tf.trainable_variables()\n",
    "                self.optimizer = tf.train.AdamOptimizer()\n",
    "                self.global_step = tf.Variable(0, name='global_step', trainable=False)\n",
    "                self.train_op = self.optimizer.minimize(self.total_loss, global_step=self.global_step)  # var_list=t_vars)\n",
    "            else:\n",
    "                self.X = tf.placeholder(tf.float32,\n",
    "                                        shape=(batch_size, 28, 28, 1))\n",
    "                self.build_arch()\n",
    "\n",
    "        tf.logging.info('Seting up the main structure')\n",
    "\n",
    "    # CapsNet 类中的build_arch方法能构建整个网络的架构\n",
    "    def build_arch(self):\n",
    "        # 以下构建第一个常规卷积层\n",
    "        with tf.variable_scope('Conv1_layer'):\n",
    "            # 第一个卷积层的输出张量为： [batch_size, 20, 20, 256]\n",
    "            # 以下卷积输入图像X,采用256个9×9的卷积核，步幅为1，且不使用\n",
    "            conv1 = tf.contrib.layers.conv2d(self.X, num_outputs=256,\n",
    "                                             kernel_size=9, stride=1,\n",
    "                                             padding='VALID')\n",
    "            # 是用 assert 可以在出现错误条件时就返回错误，有助于调整\n",
    "            assert conv1.get_shape() == [batch_size, 20, 20, 256]\n",
    "\n",
    "        # 以下是原论文中PrimaryCaps层的构建过程，该层的输出维度为 [batch_size, 1152, 8, 1]\n",
    "        with tf.variable_scope('PrimaryCaps_layer'):\n",
    "            # 调用前面定义的CapLayer函数构建第二个卷积层，该过程相当于执行八次常规卷积，\n",
    "            # 然后将各对应位置的元素组合成一个长度为8的向量，这八次常规卷积都是采用32个9×9的卷积核、步幅为2\n",
    "            primaryCaps = CapsLayer(num_outputs=32, vec_len=8, with_routing=False, layer_type='CONV')\n",
    "            caps1 = primaryCaps(conv1, kernel_size=9, stride=2)\n",
    "            assert caps1.get_shape() == [batch_size, 1152, 8, 1]\n",
    "\n",
    "        # 以下构建 DigitCaps 层, 该层返回的张量维度为 [batch_size, 10, 16, 1]\n",
    "        with tf.variable_scope('DigitCaps_layer'):\n",
    "            # DigitCaps是最后一层，它返回对应10个类别的向量（每个有16个元素），该层的构建带有Routing过程\n",
    "            digitCaps = CapsLayer(num_outputs=10, vec_len=16, with_routing=True, layer_type='FC')\n",
    "            self.caps2 = digitCaps(caps1)\n",
    "\n",
    "        # 以下构建论文图2中的解码结构，即由16维向量重构出对应类别的整个图像\n",
    "        # 1. Do masking, how:\n",
    "        with tf.variable_scope('Masking'):\n",
    "            # Method 1. masking with true label, default mode\n",
    "\n",
    "            # mask_with_y是否用真实标签蒙住目标Capsule\n",
    "            mask_with_y=True\n",
    "            if mask_with_y:\n",
    "                self.masked_v = tf.matmul(tf.squeeze(self.caps2), tf.reshape(self.Y, (-1, 10, 1)), transpose_a=True)\n",
    "                self.v_length = tf.sqrt(tf.reduce_sum(tf.square(self.caps2), axis=2, keep_dims=True) + epsilon)\n",
    "\n",
    "        # 通过3个全连接层重构MNIST图像，这三个全连接层的神经元数分别为512、1024、784\n",
    "        # [batch_size, 1, 16, 1] => [batch_size, 16] => [batch_size, 512]\n",
    "        with tf.variable_scope('Decoder'):\n",
    "            vector_j = tf.reshape(self.masked_v, shape=(batch_size, -1))\n",
    "            fc1 = tf.contrib.layers.fully_connected(vector_j, num_outputs=512)\n",
    "            assert fc1.get_shape() == [batch_size, 512]\n",
    "            fc2 = tf.contrib.layers.fully_connected(fc1, num_outputs=1024)\n",
    "            assert fc2.get_shape() == [batch_size, 1024]\n",
    "            self.decoded = tf.contrib.layers.fully_connected(fc2, num_outputs=784, activation_fn=tf.sigmoid)\n",
    "\n",
    "    # 定义 CapsNet 的损失函数，损失函数一共分为衡量 CapsNet准确度的Margin loss\n",
    "    # 和衡量重构图像准确度的 Reconstruction loss\n",
    "    def loss(self):\n",
    "        # 以下先定义重构损失，因为DigitCaps的输出向量长度就为某类别的概率，因此可以借助计算向量长度计算损失\n",
    "        # [batch_size, 10, 1, 1]\n",
    "        # max_l = max(0, m_plus-||v_c||)^2\n",
    "        max_l = tf.square(tf.maximum(0., m_plus - self.v_length))\n",
    "        # max_r = max(0, ||v_c||-m_minus)^2\n",
    "        max_r = tf.square(tf.maximum(0., self.v_length - m_minus))\n",
    "        assert max_l.get_shape() == [batch_size, 10, 1, 1]\n",
    "\n",
    "        # 将当前的维度[batch_size, 10, 1, 1] 转换为10个数字类别的one-hot编码 [batch_size, 10]\n",
    "        max_l = tf.reshape(max_l, shape=(batch_size, -1))\n",
    "        max_r = tf.reshape(max_r, shape=(batch_size, -1))\n",
    "\n",
    "        # 计算 T_c: [batch_size, 10]，其为分类的指示函数\n",
    "        # 若令T_c = Y,那么对应元素相乘就是有类别相同才会有非零输出值，T_c 和 Y 都为One-hot编码\n",
    "        T_c = self.Y\n",
    "        # [batch_size, 10], 对应元素相乘并构建最后的Margin loss 函数\n",
    "        L_c = T_c * max_l + lambda_val * (1 - T_c) * max_r\n",
    "\n",
    "        self.margin_loss = tf.reduce_mean(tf.reduce_sum(L_c, axis=1))\n",
    "\n",
    "        # 以下构建reconstruction loss函数\n",
    "        # 这一过程的损失函数通过计算FC Sigmoid层的输出像素点与原始图像像素点间的欧几里德距离而构建\n",
    "        orgin = tf.reshape(self.X, shape=(batch_size, -1))\n",
    "        squared = tf.square(self.decoded - orgin)\n",
    "        self.reconstruction_err = tf.reduce_mean(squared)\n",
    "\n",
    "        # 构建总损失函数，Hinton论文将reconstruction loss乘上0.0005\n",
    "        # 以使它不会主导训练过程中的Margin loss\n",
    "        self.total_loss = self.margin_loss + 0.0005 * self.reconstruction_err\n",
    "\n",
    "        # 以下输出TensorBoard\n",
    "        tf.summary.scalar('margin_loss', self.margin_loss)\n",
    "        tf.summary.scalar('reconstruction_loss', self.reconstruction_err)\n",
    "        tf.summary.scalar('total_loss', self.total_loss)\n",
    "        recon_img = tf.reshape(self.decoded, shape=(batch_size, 28, 28, 1))\n",
    "        tf.summary.image('reconstruction_img', recon_img)\n",
    "        self.merged_sum = tf.summary.merge_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Seting up the main structure\n",
      "INFO:tensorflow:Graph loaded\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # 训练和推断\n",
    "    capsNet = CapsNet(is_training=is_training)\n",
    "    tf.logging.info('Graph loaded')\n",
    "    sv = tf.train.Supervisor(graph=capsNet.graph,\n",
    "                             logdir=logdir,\n",
    "                             save_model_secs=0)\n",
    "\n",
    "    with sv.managed_session() as sess:\n",
    "        num_batch = int(60000 / batch_size)\n",
    "        for epoch in range(epoch):\n",
    "            if sv.should_stop():\n",
    "                break\n",
    "            for step in tqdm(range(num_batch), total=num_batch, ncols=70, leave=False, unit='b'):\n",
    "                sess.run(capsNet.train_op)\n",
    "\n",
    "            global_step = sess.run(capsNet.global_step)\n",
    "            sv.saver.save(sess, logdir + '/model_epoch_%04d_step_%02d' % (epoch, global_step))\n",
    "\n",
    "    tf.logging.info('Training done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
